import IPy
import csv
import sqlite3
import argparse
import NLAgent


csv.register_dialect('escaped', escapechar='\\', doublequote=False, quoting=csv.QUOTE_NONE, skipinitialspace=True)

def ok_IP_format(ip_string):
  if len(ip_string.split('.')) != 4:
    return False
  try:
    IPy.IP(ip_string)
    return True
  except ValueError:
    return False

def ok_IP(ip_string):
  if ok_IP_format(ip_string):
    ip, version = IPy.parseAddress(ip_string)
    if version == 4 and IPy.IP(ip_string).iptype() == 'PUBLIC':
      return True
  return False

def is_number(s):
  try:
    float(s)
    return True
  except ValueError:
    return False

def is_int(s):
  if is_number(s):
    try:
      int(s)
      return True
    except ValueError:
      return False
  return False
   
def is_port(s):
  if is_int(s):
    if int(s) < 65536 and int(s) > 0:
      return True
  return False
  
class NLAgentDB:

  def add_host(self, host, delay, drops):
    """Add info about a host. Delays are in microseconds."""
    host_id = self.get_host_id(host)
    if host_id != None:
      print("Skipping adding host since we already have", host, "in the database.")
      return host_id
    else:
      if ok_IP(host):
        if is_number(delay) and is_number(drops):
          self.cur.execute('INSERT INTO hosts (host_ip, delay, drops) VALUES(?,?,?)',(host,delay,drops)) 
          return self.cur.lastrowid
        else:
          print("Skipping ",host," due to non-number delay or drops:",delay," ",drop)
          return None
      else:
        print("Skipping",host,"due to invalid IP")
        return None
    return None

  def get_host_id(self, host):
    #TABLE hosts (host_id INTEGER PRIMARY KEY, host_ip text, delay real, drops integer)''')
    value = (host)
    for row in self.cur.execute('SELECT host_id FROM hosts WHERE host_ip=:ip', {"ip": host}):
      return row[0]
    return None
    
  def get_host_delay(self, host):
    for row in self.cur.execute('SELECT delay FROM hosts WHERE host_ip=?',(host,)):
      return row[0]
    return None
    
  def add_connection(self, src, src_port, dst, dst_port, s_time, conn_id):
    if ok_IP(src) and ok_IP(dst): 
      if is_port(src_port) and is_port(dst_port):
        self.cur.execute('INSERT INTO connections (src,dst,src_port,dst_port,start_time,conn_id) VALUES (?,?,?,?,?,?)',(src,dst,src_port,dst_port,s_time,conn_id))
        return conn_id
      else:
        print("WARNING: Skipping due to bad port number", src,":",src_port,"->",dst,":",dst_port)
        return None
    else:
      print("WARNING: Skipping", src, "->", dst, "due to invalid IP")
      return None
    return None
  
  def add_event(self,conn_id,seq_num,actor_ip,state,value):
      event_id = self.get_event_id(conn_id.replace(' ', ''));
      if (event_id != None):
        if state == 'WAIT':
        # Put the sleep in milliseconds
          value = int(float(value) * 1000)
        self.cur.execute('INSERT INTO events (conn_id,seq_num,actor_ip,state,value) VALUES (?,?,?,?,?)',(conn_id,seq_num,actor_ip,state,value))
        return self.cur.lastrowid
      return None
    #events (event_id KEY, conn_id integer, seq_num, actor state text, value integer, FOREIGN KEY(conn_id) REFERENCES events(event_id))''')
  
  def get_event_id(self,conn_id):
    for row in self.cur.execute('SELECT conn_id FROM connections WHERE conn_id=?',(conn_id,)):
      return row[0]    
    return None
  
  def get_connection_id(self, src, src_port, dst, dst_port, start_time):
    #connections (conn_id, src, dst, src_port, dst_port, start_time)
    for row in self.cur.execute('SELECT conn_id, start_time FROM connections WHERE src=? AND src_port=? AND dst=? AND dst_port=?',(src,src_port,dst,dst_port)):
      return row[0]
  
  def get_connection(self,conn_id):
    if is_int(conn_id):
      for row in self.cur.execute('SELECT src, src_port, dst, dst_port FROM connections WHERE conn_id=?', (conn_id,)):
        return {'src':row[0], 'src_port':row[1], 'dst':row[2], 'dst_port':row[3]}
    else:
      print("ERROR: Look up for connection ID:",conn_id,"failed")
      return None
  
  def get_last_connection_time(self):
    for row in self.cur.execute('SELECT start_time, conn_id, src FROM connections ORDER BY start_time DESC limit 1'):
      return row
    return None
    
  
  def get_list_of_connections(self):
    r = self.cur.execute('SELECT start_time, conn_id, src, src_port, dst, dst_port FROM connections')
    return r
  
  def get_connection_start_time(self, conn_id):
    for row in self.cur.execute('SELECT start_time FROM connections WHERE conn_id=?',(conn_id,)):
      return row[0]
    print("ERROR: Look up starttime for connection ID:",conn_id,"failed")
    return None
    
  def get_event(self, conn_id, seq_num):
    for row in self.cur.execute('SELECT seq_num, actor_ip,state,value from events WHERE conn_id=? AND seq_num=?',(conn_id, seq_num)):
      return {'seq_num': row[0], 'actor_ip':row[1], 'state':row[2], 'value':row[3]}
      #return seq_num, actor_ip, state, value
    return None
    
  def get_event_list(self,conn_id):
    list = self.cur.execute('SELECT seq_num, actor_ip, state, value from events WHERE conn_id=? ORDER BY seq_num',(conn_id,))
    return list

  def get_all_events(self):
    list = self.cur.execute('SELECT conn_id, seq_num, actor_ip, state, value from events ORDER BY conn_id, seq_num')
    return list

  def next_event_exists(self, conn_id, last_seq_num):
    for row in self.cur.execute('SELECT seq_num from events WHERE conn_id=? AND seq_num=?',(conn_id, last_seq_num+1)):
      return row[0]
    return None

  def reduce_tables(self, hostfile=None):
    """There are a number of ways tables may have excess information.
     Here we cull the info"""
    # The main excess of information is in the event database.
    # Merge data transfer events
    # Add in wait times between events
    # Add DONE event if needed (ie. we didn't see a FIN)
    
    # Once the event database is reduced, we can remove
    # connections (from both the event and connection database)
    # which do not have bidirectional communication.
    
    # If we're given a hostfile, we can reduce even further
    # since this tells us the hosts we care about, and
    # we can get rid of information on the rest.


  def add_hosts_from_csv(self,filename):
    try:
      with open(filename) as f: pass
    except IOError as e:
      print("Skipping adding to host table from ", filename, ":", e)
      
    self.cur.execute("DELETE FROM hosts")
    
    hostCSVReader = csv.reader(open(filename, 'r'), delimiter=',', dialect='escaped')
    
    f = open(filename, 'rt')
    try:
      reader = csv.DictReader(f)
      for row in reader:
        if len(row) == 5:
          self.add_host(row['HOST'],row['DELAY'],row['DROP'])
        else:
          print("Poorly formatted line:", row, "Skipping.")
    finally:
      f.close()

  def add_conns_from_csv(self,filename):
    try:
      with open(filename) as f: pass
    except IOError as e:
      print("Skipping adding to conn table from ", filename, ":", e)

    self.cur.execute("DELETE FROM connections")    
    
    hostCSVReader = csv.reader(open(filename, 'r'), delimiter=',', dialect='escaped')
    
    f = open(filename, 'rt')
    try:
      reader = csv.DictReader(f)
      for row in reader:
        if len(row) == 8:
          self.add_connection(row['SRC'], row['SPORT'], row['DST'], row['DPORT'], row['STIME'], row['CONNID'])
        else:
          print("Poorly formatted line:", row, "Skipping.")
    finally:
      f.close()
    self.conn.commit()
      
  def add_events_from_csv(self,filename):
    try:
      with open(filename) as f: pass
    except IOError as e:
      print("Skipping adding to conn table from ", filename, ":", e)
    
    self.cur.execute("DELETE FROM events")
    eventCSVReader = csv.reader(open(filename, 'r'), delimiter=',', dialect='escaped')
    
    f = open(filename, 'rt')
    try:
      reader = csv.DictReader(f)
      for row in reader:
        if len(row) == 6:
          #X,CONNID,SEQ_NUM,ACTOR,STATE,VALUE
          self.add_event(row['CONNID'],row['SEQ_NUM'],row['ACTOR'],row['STATE'],row['VALUE'])
        else:
          print("Poorly formatted line:", row, "Skipping.")
    finally:
      f.close()
    self.conn.commit()
    
  
    
  def __init__(self,filename):
    
    
    self.filename = filename
    create_db = False
    try:
      with open(self.filename) as f: pass
    except IOError as e:
      print("Will need to create database. '", filename,"' does not yet exist.")
      create_db = True
    self.conn = sqlite3.connect(filename)
    self.conn.row_factory = sqlite3.Row
    self.cur = self.conn.cursor()
    if create_db == True:
      self.__create_tables__()
  
  def close(self):
    self.cur.close()
  
  def __create_tables__(self):
    self.cur.execute('''CREATE TABLE hosts (host_id INTEGER PRIMARY KEY, host_ip text, delay real, drops integer)''')
    self.cur.execute('''CREATE TABLE connections (conn_id INTEGER PRIMARY KEY, src text, dst text, src_port integer, dst_port integer, start_time float)''')
    self.cur.execute('''CREATE TABLE events (event_id INTEGER PRIMARY KEY, conn_id integer, seq_num integer, state text, actor_ip text, value integer, FOREIGN KEY(conn_id) REFERENCES events(event_id))''')
  
  
if __name__ == '__main__':
  parser = argparse.ArgumentParser(parents=[NLAgent.parser])
  db_args = parser.add_argument_group('Database Options')
  db_args.add_argument('--host_csv_file', action="store", dest='host_csv')
  db_args.add_argument('--connection_csv_file', action="store", dest='conn_csv')
  db_args.add_argument('--event_csv_file', action="store", dest='event_csv')
  db_args.add_argument('--reduce_db', nargs='?', const='Reduce Current DB', default=None)
  args = parser.parse_args()

  database = NLAgentDB(args.filename)
  
  if args.host_csv != None:
    print("Reading from '", args.host_csv, "'.")
    database.add_hosts_from_csv(args.host_csv)
  if args.conn_csv != None:
    print("Reading from '",args.conn_csv,"'.")
    database.add_conns_from_csv(args.conn_csv)
  if args.event_csv != None:
    print("Reading from '",args.event_csv, "'.")
    database.add_events_from_csv(args.event_csv)

  if args.reduce_db == None:
    print("You don't want me to reduce the tables.")
  elif args.reduce_db == 'Reduce Current DB':
    print("You want me to reduce the tables in current database.")
  else:
    print("You want me to reduce the tables and put them in", args.reduce_db)
    
  database.conn.commit()
  database.cur.close()
    
